# -*- coding: utf-8 -*-
"""
交易日历工具
用于判断是否为交易日，获取交易日列表等
"""
import pandas as pd
import akshare as ak
from datetime import datetime, timedelta
from typing import List, Optional
import logging

logger = logging.getLogger(__name__)

class TradingCalendar:
    """交易日历类"""
    
    def __init__(self):
        self._trading_days_cache = None
        self._cache_start_date = None
        self._cache_end_date = None
    
    def is_trading_day(self, date: datetime) -> bool:
        """判断指定日期是否为交易日"""
        if isinstance(date, str):
            date = datetime.strptime(date, '%Y-%m-%d')
        
        # 检查是否为周末
        if date.weekday() >= 5:  # 周六(5)或周日(6)
            return False
        
        # 检查是否为交易日（通过缓存或实时查询）
        trading_days = self.get_trading_days(
            start_date=date.strftime('%Y-%m-%d'),
            end_date=date.strftime('%Y-%m-%d')
        )
        
        return len(trading_days) > 0
    
    def get_trading_days(self, start_date: str, end_date: str) -> List[str]:
        """获取指定日期范围内的交易日列表"""
        try:
            # 如果缓存覆盖了查询范围，直接使用缓存
            if (self._trading_days_cache and 
                self._cache_start_date and 
                self._cache_end_date and
                start_date >= self._cache_start_date and 
                end_date <= self._cache_end_date):
                
                trading_days = [day for day in self._trading_days_cache 
                              if start_date <= day <= end_date]
                return trading_days
            
            # 从akshare获取交易日历
            logger.info(f"正在获取交易日历: {start_date} - {end_date}")
            df = ak.tool_trade_date_hist_sina()
            
            if df.empty:
                logger.warning("未获取到交易日历数据")
                return []
            
            # 处理日期格式
            df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y-%m-%d')
            
            # 过滤指定日期范围
            trading_days = df[
                (df['trade_date'] >= start_date) & 
                (df['trade_date'] <= end_date)
            ]['trade_date'].tolist()
            
            # 更新缓存
            self._trading_days_cache = df['trade_date'].tolist()
            self._cache_start_date = df['trade_date'].min()
            self._cache_end_date = df['trade_date'].max()
            
            logger.info(f"获取到 {len(trading_days)} 个交易日")
            return trading_days
            
        except Exception as e:
            logger.error(f"获取交易日历失败: {e}")
            # 如果获取失败，使用简单的周末判断
            return self._get_weekdays_fallback(start_date, end_date)
    
    def _get_weekdays_fallback(self, start_date: str, end_date: str) -> List[str]:
        """备用方法：只排除周末，不排除节假日"""
        start = datetime.strptime(start_date, '%Y-%m-%d')
        end = datetime.strptime(end_date, '%Y-%m-%d')
        
        trading_days = []
        current = start
        while current <= end:
            if current.weekday() < 5:  # 周一到周五
                trading_days.append(current.strftime('%Y-%m-%d'))
            current += timedelta(days=1)
        
        logger.warning(f"使用备用方法获取交易日，共 {len(trading_days)} 天")
        return trading_days
    
    def get_next_trading_day(self, date: str) -> Optional[str]:
        """获取指定日期后的下一个交易日"""
        if isinstance(date, str):
            date = datetime.strptime(date, '%Y-%m-%d')
        
        # 从当前日期的下一天开始查找
        current = date + timedelta(days=1)
        
        # 最多查找30天
        for _ in range(30):
            if self.is_trading_day(current):
                return current.strftime('%Y-%m-%d')
            current += timedelta(days=1)
        
        return None
    
    def get_previous_trading_day(self, date: str) -> Optional[str]:
        """获取指定日期前的上一个交易日"""
        if isinstance(date, str):
            date = datetime.strptime(date, '%Y-%m-%d')
        
        # 从当前日期的前一天开始查找
        current = date - timedelta(days=1)
        
        # 最多查找30天
        for _ in range(30):
            if self.is_trading_day(current):
                return current.strftime('%Y-%m-%d')
            current -= timedelta(days=1)
        
        return None
    
    def get_trading_days_between(self, start_date: str, end_date: str) -> List[str]:
        """获取两个日期之间的所有交易日"""
        return self.get_trading_days(start_date, end_date)
    
    def get_missing_trading_days(self, stock_code: str, db_manager) -> List[str]:
        """获取指定股票缺失的交易日列表"""
        try:
            # 获取股票的最新交易日期
            sql = "SELECT MAX(trade_date) as latest_date FROM new_stock_daily WHERE stock_code = %s"
            result = db_manager.execute_query(sql, (stock_code,))
            
            if not result or not result[0]['latest_date']:
                # 如果没有历史数据，从一年前开始
                start_date = (datetime.now() - timedelta(days=365)).strftime('%Y-%m-%d')
            else:
                # 从最新日期的下一天开始
                latest_date = result[0]['latest_date']
                next_day = latest_date + timedelta(days=1)
                start_date = next_day.strftime('%Y-%m-%d')
            
            # 到今天为止
            end_date = datetime.now().strftime('%Y-%m-%d')
            
            # 获取交易日列表
            trading_days = self.get_trading_days(start_date, end_date)
            
            # 过滤掉已经同步的日期
            if trading_days:
                placeholders = ','.join(['%s'] * len(trading_days))
                sql = f"""
                SELECT trade_date FROM new_stock_daily 
                WHERE stock_code = %s AND trade_date IN ({placeholders})
                """
                synced_dates = db_manager.execute_query(sql, (stock_code, *trading_days))
                synced_date_list = [row['trade_date'].strftime('%Y-%m-%d') for row in synced_dates]
                
                missing_days = [day for day in trading_days if day not in synced_date_list]
                return missing_days
            
            return []
            
        except Exception as e:
            logger.error(f"获取股票 {stock_code} 缺失交易日失败: {e}")
            return []

def main():
    """测试交易日历功能"""
    calendar = TradingCalendar()
    
    # 测试当前日期是否为交易日
    today = datetime.now().strftime('%Y-%m-%d')
    print(f"今天 {today} 是否为交易日: {calendar.is_trading_day(today)}")
    
    # 测试获取最近一周的交易日
    end_date = datetime.now().strftime('%Y-%m-%d')
    start_date = (datetime.now() - timedelta(days=7)).strftime('%Y-%m-%d')
    trading_days = calendar.get_trading_days(start_date, end_date)
    print(f"最近一周的交易日: {trading_days}")
    
    # 测试获取下一个交易日
    next_trading_day = calendar.get_next_trading_day(today)
    print(f"下一个交易日: {next_trading_day}")

if __name__ == "__main__":
    main()
